package com.sparrow.socket.tcp.handler;

import com.sparrow.common.util.CollectionUtil;
import com.sparrow.common.util.IpUtil;
import com.sparrow.common.util.PropertiesUtil;
import io.netty.channel.ChannelHandlerContext;
import io.netty.channel.ChannelInboundHandlerAdapter;
import io.netty.util.AttributeKey;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import java.net.InetSocketAddress;
import java.net.SocketAddress;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CopyOnWriteArrayList;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * ip连接数限制 + ip白名单 + ip黑名单
 */
public class IpLimitHandler extends ChannelInboundHandlerAdapter {

    private static Logger logger = LoggerFactory.getLogger(IpLimitHandler.class);

    private static int MAX_CONNECTIONS_PER_IP = PropertiesUtil.getPropertyAsIntOr("ip_conn_limit", 5);

    // key: IP 地址，value: 当前连接数
    private final ConcurrentHashMap<String, AtomicInteger> ipConnectionCount = new ConcurrentHashMap<>();

    // 每个 Channel 保存它的 IP，便于关闭时清理
    private static final AttributeKey<String> ATTR_IP = AttributeKey.valueOf("ip");

    // 黑名单 IP：禁止连接
    private static final CopyOnWriteArrayList<String> blackList = new CopyOnWriteArrayList<>(CollectionUtil.toStringSet(PropertiesUtil.getProperty("ip_black_list")));

    // 白名单 IP：不受连接数限制
    private static final CopyOnWriteArrayList<String> whiteList = new CopyOnWriteArrayList<>(CollectionUtil.toStringSet(PropertiesUtil.getProperty("ip_white_list")));

    @Override
    public void channelActive(ChannelHandlerContext ctx) throws Exception {
        String ip = getIp(ctx);
        ctx.channel().attr(ATTR_IP).set(ip);
        if (!blackList.isEmpty()) {
            for (String cidr : blackList) {
                if (IpUtil.isInRange(ip, cidr)) {
                    logger.error("IP{} 在黑名单中，拒绝连接", ip);
                    ctx.close();
                    return;
                }
            }
        }
        if (!whiteList.isEmpty()) {
            for (String cidr : whiteList) {
                if (IpUtil.isInRange(ip, cidr)) {
                    super.channelActive(ctx);
                    return;
                }
            }
        }
        AtomicInteger counter = ipConnectionCount.computeIfAbsent(ip, k -> new AtomicInteger(0));
        int currentCount = counter.incrementAndGet();
        if (currentCount > MAX_CONNECTIONS_PER_IP) {
            // 超过连接数限制
            logger.error("IP{} 超过最大连接数限制，当前{}", ip, currentCount);
            counter.decrementAndGet();
            ctx.close(); // 关闭连接
            return;
        }
        super.channelActive(ctx);
    }

    @Override
    public void channelInactive(ChannelHandlerContext ctx) throws Exception {
        String ip = ctx.channel().attr(ATTR_IP).get();
        if (ip != null) {
            AtomicInteger counter = ipConnectionCount.get(ip);
            if (counter != null) {
                int remaining = counter.decrementAndGet();
                if (remaining <= 0) {
                    ipConnectionCount.remove(ip);
                }
            }
        }
        super.channelInactive(ctx);
    }

    private String getIp(ChannelHandlerContext ctx) {
        SocketAddress remoteAddr = ctx.channel().remoteAddress();
        if (remoteAddr instanceof InetSocketAddress) {
            return ((InetSocketAddress) remoteAddr).getAddress().getHostAddress();
        }
        return "unknown";
    }

}
